#pragma once

// #include "lcir/IR/Argument.h"
#include <lcir/ADT/IList.h>
#include <lcir/ADT/OwningRefVector.h>
#include <lcir/IR/Context.h>
#include <lcir/IR/Expression.h>
#include <lcir/IR/LambdaBase.h>
#include <lcir/IR/Prototype.h>
#include <lcir/IR/Type.h>
#include <lcir/IR/Value.h>

#include <functional>
#include <iterator>
#include <memory>
#include <vector>
namespace lcir {
// 处理walk相关的代码
template <bool IsConst>
struct WalkTrait;
template <>
struct WalkTrait<true> {
  using lambdaType = const Lambda *;
};
template <>
struct WalkTrait<false> {
  using lambdaType = Lambda *;
};
class LambdaProto;
class LambdaType;
enum WalkResult {
  Interrupt,
  Advance,
  Skip,  // TODO...
};
inline WalkResult walkInterrupt() { return Interrupt; }
inline WalkResult walkAdvance() { return Advance; }
inline WalkResult walkSkip() { return Skip; }


class Lambda final : public ValueCRTP<Lambda>, public LambdaBase {
  friend class Expression;
  // std::unique_ptr<LambdaProto> proto;
  std::string name = "lambda";
  // 注意：lambda被定义的expr中，lambda是作为操作数而不是result出现！
  Expression *defExpr{nullptr};
  OwningRefVector<Argument> arguments;

 protected:
  /**
   * @param: 设置新的proto，主要给lambda添加参数/重新设置参数时候使用
   */
  void setProto(std::unique_ptr<LambdaProto> proto);

 public:
  using exprList = IList<Expression>;
  Lambda(std::unique_ptr<LambdaProto> proto);
  Lambda();
  Lambda(const std::vector<Type *> &paramTypes, Type *resultType);

  // lambda的一些静态方法，由context管理lambda生命周期
  static Lambda *get(const std::vector<Type *> &paramTypes, Type *resultType) {
    auto *ctx = GlobalContext::get();
    return ctx->getLamda(paramTypes, resultType);
  }

  static Lambda *get() { return get({}, nullptr); }

  static Lambda *get(const std::vector<Type *> &paramTypes) {
    auto *ctx = GlobalContext::get();
    return ctx->getLamda(paramTypes, nullptr);
  }
  Expression *getDefExpr() const override { return defExpr; }

  LambdaProto *getLambdaProto() const;
  LambdaType *getLambdaType() const;

  std::vector<Type *> getParamTypes() const;
  Type *getResultType() const;

  std::vector<Argument *> getArgs() const;
  Argument *getArg(size_t idx) const;

  std::string print() const;

  /**
   * @brief: 添加参数，注意：添加lambda参数后，lambda的类型会改变！
   */
  void addArg(Type *type);

  /// TODO: 这俩暂时用不到，以后再加
  // void addArgs(const std::vector<Type *> &typs);
  // void setArgs(const std::vector<Type *> &typs);

  static ValueKind getValueKind() { return ValueKind::VK_Lambda; }

  // 获得当前lambda节点的嵌套深度
  size_t getDepth() const;

  /**
   * @brief: 把区间内的expression从链表里卸载出来
   */
  [[nodiscard]] std::vector<std::unique_ptr<Expression>> removeFromLambda(Expression *begin, Expression *end);

  /**
   * @brief:  把lambda内的expression卸载出来，创建一个新的lambda
   */
  Lambda *makeNewLambda();

  /**
   * @brief: 判断该Lambda是否已经密封
   */
  bool isSealed() const;

  template <bool IsBackWard = false>
  void walk(const std::function<void(const Expression *expr)> &callback) const {
    walkImpl<IsBackWard, true>(callback);
  }
  template <bool IsBackWard = false>
  void walk(const std::function<void(Expression *expr)> &callback) {
    walkImpl<IsBackWard, false>(callback);
  }
  template <class OpTy, bool IsBackWard = false, class Fn>
  void walkOnOp(const Fn &callback) const {
    walkImpl<OpTy, IsBackWard, true>(callback);
  }
  template <class OpTy, bool IsBackWard = false, class Fn>
  void walkOnOp(const Fn &callback) {
    walkImpl<OpTy, IsBackWard, false>(callback);
  }

 private:
  void setDefExpr(Expression *expr);
  void initArgs(const std::vector<Type *> &argTypes);

  /*-------------walk方法的模板函数-------------*/
  template <bool IsBackWard, bool IsConst, class Fn>
  void walkImpl(const Fn &callback) const {
    if constexpr (IsBackWard) {
      auto range = ReverseRange(getExprs());
      auto incrange = IncRange(range);
      walkBase<IsBackWard, IsConst>(callback, incrange);
    } else {
      auto incrange = IncRange(getExprs());
      walkBase<IsBackWard, IsConst>(callback, incrange);
    }
  }
  template <bool IsBackWard, bool IsConst, class Fn>
  void walkImpl(const Fn &callback) {
    if constexpr (IsBackWard) {
      auto range = ReverseRange(getExprs());
      auto incrange = IncRange(range);
      walkBase<IsBackWard, IsConst>(callback, incrange);
    } else {
      auto incrange = IncRange(getExprs());
      walkBase<IsBackWard, IsConst>(callback, incrange);
    }
  }
  template <class OpTy, bool IsBackWard, bool IsConst, class Fn>
  void walkImpl(const Fn &callback) const {
    if constexpr (IsBackWard) {
      auto range = ReverseRange(getExprs());
      auto incrange = IncRange(range);
      walkBase<IsBackWard, IsConst>(callback, incrange);
    } else {
      auto incrange = IncRange(getExprs());
      walkBase<OpTy, IsBackWard, IsConst>(callback, incrange);
    }
  }

  template <class OpTy, bool IsBackWard, bool IsConst, class Fn>
  void walkImpl(const Fn &callback) {
    if constexpr (IsBackWard) {
      auto range = ReverseRange(getExprs());
      auto incrange = IncRange(range);
      walkBase<IsBackWard, IsConst>(callback, incrange);
    } else {
      auto incrange = IncRange(getExprs());
      walkBase<OpTy, IsBackWard, IsConst>(callback, incrange);
    }
  }

  template <bool IsBackWard, bool IsConst, class RangeTy, class Fn>
  void walkBase(const Fn &callback, RangeTy &&range) const {
    using lambdaType = typename WalkTrait<IsConst>::lambdaType;
    for (auto &expr : range) {
      callback(&expr);
      for (auto *oprd : expr.getOperands()) {
        if (lambdaType lambda = oprd->getValue()->template dyn_cast<Lambda>())
          lambda->template walk<IsBackWard>(callback);
      }
    }
  }
  template <class OpTy, bool IsBackWard, bool IsConst, class RangeTy, class Fn>
  void walkBase(const Fn &callback, RangeTy &&range) const {
    using lambdaType = typename WalkTrait<IsConst>::lambdaType;
    bool skip{false};
    for (auto &expr : range) {
      if (!expr.template isOp<OpTy>()) continue;
      if constexpr (std::is_void_v<std::invoke_result_t<Fn, Expression *>>) {
        callback(&expr);
      } else if constexpr (std::is_same_v<WalkResult, std::invoke_result_t<Fn, Expression *>>) {
        WalkResult result = callback(&expr);
        if (result == Interrupt) return;
        if (result == Skip) skip = true;
      }
      if (!skip)
        for (auto *oprd : expr.getOperands()) {
          if (lambdaType lambda = oprd->getValue()->template dyn_cast<Lambda>())
            lambda->template walkOnOp<OpTy, IsBackWard>(callback);
        }
    }
  }

 public:
  Lambda(Lambda &&other) = delete;
  Lambda &operator=(Lambda &&other) = delete;
  Lambda(const Lambda &other) = delete;
  Lambda &operator=(const Lambda &other) = delete;
  ~Lambda() = default;
};

}  // namespace lcir
